import models_mae
from util.LoRALayer import LoRALayer
import torch.nn as nn


class MAEWithLoRA(models_mae.MaskedAutoencoderViT):
    """
    在 MAE (Masked Autoencoder) 模型的基础上，替换全连接层为 LoRA (Low-Rank Adaptation) 层。
    该模型继承自 `MaskedAutoencoderViT` 类，并在其基础上添加了 LoRA 层。

    参数：
        *args, **kwargs: 传递给父类 `MaskedAutoencoderViT` 的其他参数。
        r (int): LoRA 层的秩（rank），控制适配器的参数数量。
        alpha (float): LoRA 层的缩放因子，用于调整适配器的影响。
    """

    def __init__(self, *args, r=4, alpha=1.0, **kwargs):
        # 初始化父类 MaskedAutoencoderViT
        super(MAEWithLoRA, self).__init__(*args, **kwargs)

        # 替换模型中的全连接层（假设模型有一层 self.fc）为 LoRA 层
        # self.fc 是原始模型中的全连接层，LoRALayer 是自定义的 LoRA 适配器层
        self.lora_fc = LoRALayer(self.fc.in_features, self.fc.out_features, r=r, alpha=alpha)

        # 如果不希望使用原始的全连接层，可以将其置为空（可选）
        # self.fc = nn.Identity()  # 将原来的全连接层替换为一个身份层，防止其被使用

    def forward(self, x):
        """
        前向传播函数，替换了原始模型的全连接层为 LoRA 层。

        参数：
            x (Tensor): 输入的张量，通常是经过掩码处理的图像或特征。

        返回：
            Tensor: 输出的张量，经过 LoRA 层处理后的结果。
        """
        # 调用父类的 forward 函数进行前向传播
        x = super(MAEWithLoRA, self).forward(x)

        # 使用 LoRA 层替代原始全连接层
        x = self.lora_fc(x)  # 使用 LoRA 层进行特征变换

        # 返回经过 LoRA 适配后的输出
        return x

